{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1.2 CoefficientFunctions\n",
    "====\n",
    "\n",
    "In NGSolve, `CoefficientFunction`s are representations of functions defined on the computational domain. Examples are expressions of coordinate variables $x, y, z$ and functions that are  constant on subdomains. Much of the magic behind the seamless integration of NGSolve with python lies in `CoefficientFunction`s. This tutorial introduces you to them.\n",
    "\n",
    "After this tutorial you will know how to \n",
    "\n",
    "- **define** a `CoefficientFunction`,\n",
    "- **visualize**  a `CoefficientFunction`,\n",
    "- **evaluate** `CoefficientFunction`s at points,\n",
    "- print the **expression tree** of `CoefficientFunction`,\n",
    "- **integrate** a `CoefficientFunction`, \n",
    "- **differentiate** a `CoefficientFunction`, \n",
    "- include **parameter** in `CoefficientFunction`s,\n",
    "- **interpolate** a `CoefficientFunction` into a finite element space, \n",
    "- define **vector-valued** `CoefficientFunction`s, and \n",
    "- **compile** `CoefficientFunction`s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "from netgen.geom2d import unit_square\n",
    "import matplotlib.pyplot as plt\n",
    "mesh = Mesh (unit_square.GenerateMesh(maxh=0.2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "myfunc = x*(1-x)\n",
    "myfunc   # You have just created a CoefficientFunction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x        # This is a built-in CoefficientFunction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the function\n",
    "\n",
    "Use the `mesh` to visualize the function in Netgen's GUI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw(myfunc, mesh, \"firstfun\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate the function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mip = mesh(0.2, 0.2)\n",
    "myfunc(mip)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `myfunc(0.2,0.3)` will not evaluate the function: one must give points in the form of `MappedIntegrationPoint`s like `mip` above. The `mesh` knows how to produce them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examining functions on sets of points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pts = [(0.1*i, 0.2) for i in range(10)]\n",
    "vals = [myfunc(mesh(*p)) for p in pts] \n",
    "for p,v in zip(pts, vals):\n",
    "    print(\"point=(%3.2f,%3.2f), value=%6.5f\"\n",
    "         %(p[0], p[1], v))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We may plot the restriction of the `CoefficientFunction` on a line using matplotlib. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "px = [0.01*i for i in range(100)]\n",
    "vals = [myfunc(mesh(p,0.5)) for p in px]\n",
    "plt.plot(px,vals)\n",
    "plt.xlabel('x')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Expression tree of a function\n",
    "\n",
    "Internally, coefficient functions are implemented as an expression tree made from building blocks like `x`, `y`, `sin`, etc., and arithmetic operations.\n",
    "\n",
    "E.g., the expression tree for `myfunc = x*(1-x)` looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(myfunc) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integrate a function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can numerically integrate the function using the mesh:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Integrate(myfunc, mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can change the precision of the quadrature rule used for the integration using the key word argument `order`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Differentiate a function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Automatic differentiation of a `CoefficientFunction` is now possible through the `Diff` method. Here is how you get $\\partial / \\partial x$ of `myfunc`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_myfunc = myfunc.Diff(x)\n",
    "Draw(diff_myfunc, mesh, \"derivative\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See if you can recognize an implementation of the product rule \n",
    "\n",
    "$$\n",
    "\\frac{\\partial}{\\partial x} x (1-x)\n",
    "= \n",
    "\\frac{\\partial x}{\\partial x} \n",
    "(1-x) + \n",
    "x\\frac{\\partial (1-x)}{\\partial x} \n",
    "$$\n",
    "\n",
    "in the tree-representation of the differentiated coefficient function, printed below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(diff_myfunc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters in functions\n",
    "\n",
    "When building complex coefficient functions from simple ones like `x` and `y`, you may often want to introduce `Parameter`s, which are constants whose value may be changed later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = Parameter(1.0)\n",
    "f = sin(k*y)\n",
    "Draw(f, mesh, \"f\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same `f` may be given a different value of `k` later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k.Set(10)\n",
    "Draw(f, mesh, \"f\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the expression tree of `f`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how the `Parameter` \n",
    "is now a **node** in the expression tree.  You can differentiate a coefficient function with respect to such quantities by passing it as argument to `Diff`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f.Diff(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Integrate((f.Diff(k) - y*cos(k*y))**2, mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolate a function\n",
    "\n",
    "We may `Set` a `GridFunction` using a `CoefficientFunction`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fes = H1(mesh, order=1)\n",
    "u = GridFunction(fes)\n",
    "u.Set(myfunc)\n",
    "Draw(u)         # Cf. Draw(myfunc, mesh, \"firstfun\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* The `Set` method interpolates `myfunc` to an element `u` in the finite element space.\n",
    "\n",
    "* `Set` does an *Oswald-type interpolation* as follows:\n",
    "    - It first zeros the grid function;\n",
    "    - It then projects `myfunc` in $L^2$ on each mesh element;\n",
    "    - It then averages dofs on element interfaces for conformity.\n",
    "    \n",
    "* We will see other ways to interpolate in [2.10](../unit-2.10-dualbasis/dualbasis.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vector-valued `CoefficientFunction`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an example of a vector-valued coefficient function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vecfun = CoefficientFunction((-y, sin(x)))\n",
    "Draw(vecfun, mesh, \"vecfun\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Click on `Draw Surface Vectors` in the `Visual` menu to visualize this vector field.\n",
    "\n",
    "Another example of a vector-valued coefficient function is the gradient of the above-set `GridFunction`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u.Set(myfunc)\n",
    "gradu = grad(u)\n",
    "Draw(gradu, mesh, \"grad_firstfun\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compiled `CoefficientFunction`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation of a `CoefficientFunction` at a point is usually done \n",
    "by traversing its expression tree and evaluating each node of the tree. When the tree has repeated nodes, this is likely wasteful. NGSolve allows you to **\"compile\"** a `CoefficientFunction` to increase the efficiency of its evaluation. The compilation translates the expression tree into a sequence of linear steps. \n",
    "\n",
    "Continuing with our simple `myfunc` example, here is how to use the `Compile` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "myfunc_compiled = myfunc.Compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now look at the differences between the compiled and non-compiled `CoefficientFunction`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(myfunc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(myfunc_compiled)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation of the compiled function is now a linear sequence of Steps 0, 1, 2, and 3 above. We understand the printed description of Steps 2 and 3 above to mean the following.\n",
    "\n",
    "*  Step 2: (Output of Step 1) - (Output of Step 0) \n",
    "*  Step 3: (Output of Step 0) * (Output of Step 2) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is another example, along with differences in timings for integrating the coefficient function on the mesh. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f0 = myfunc\n",
    "f1 = f0*y \n",
    "f2 = f1*f1 + f1*f0 + f0*f0 \n",
    "f3 = f2*f2*f2*f0**2 + f0*f2**2 + f0**2 + f1**2 + f2**2\n",
    "final = f3 + f3 + f3 \n",
    "finalc = final.Compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit Integrate(final, mesh, order=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit Integrate(finalc, mesh, order=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If your NGSolve installation has a compiler script (you likely do if you built from source using a compiler), then you have the option of letting that compiler optimize your coefficient function further. Here is an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "finalcc = final.Compile(realcompile=True, wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit Integrate(finalcc, mesh, order=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(finalc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(final)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
